#pragma once

#include <array>
#include <functional>
#include <string>

#include "backend/device/txda_device.h"
#include "backend/kernel_interface/kernel_args.h"
#include "backend/program/program_mainc_group.h"


namespace txdnn {
namespace backend {

/* struct KernelInvokeMainCGroup */
struct KernelInvokeMainCGroup {
public:
  KernelInvokeMainCGroup() {}
  KernelInvokeMainCGroup(txdaStream_t stream,
                         txdaFunction_t fn,
                         std::array<size_t, 3> pldims,
                         std::array<size_t, 3> pgdims,
                         std::string name,
                         std::function<void(txdaEvent_t, txdaEvent_t)> callback,
                         bool multi_graph_launch) 
    : stream_(stream),
      fn_(fn),
      ldims_(pldims),
      gdims_(pgdims),
      name_(name),
      callback_(callback),
      multi_graph_launch_(multi_graph_launch) {}

  void operator()(std::vector<OpKernelArg>& args) const {
    // TODO memory_copy and kernel_launch
    // call run()
  }

  template <class... Ts>
  void operator()(Ts... xs) const {
    // TODO call run() or run_multi_graph()

    if (multi_graph_launch_) {
      auto args = std::array<void*, sizeof...(xs)>{(&xs)...};
      run_multi_graph(args.data());
    } else {
      // TODO
    }
  }

  const std::string& getName() const { return name_; }

private:
  void run(void* args, std::size_t size) const {
    // TODO kernel_launch
    printf("run() : Kernel launched\n");
  }

  void run_multi_graph(void** kernel_args) const {
    // TODO kernel_launch
    printf("run_multi_graph() : Kernel launched\n");
  }

  std::string name_;
  txdaStream_t stream_          = nullptr;
  txdaFunction_t fn_            = nullptr;
  std::array<size_t, 3> ldims_  = {};
  std::array<size_t, 3> gdims_  = {};
  std::function<void(txdaEvent_t, txdaEvent_t)> callback_;
  bool multi_graph_launch_;
};


/* struct KernelMainCGroup */
struct KernelMainCGroup {
public:
  KernelMainCGroup() {}
  KernelMainCGroup(ProgramMainCGroup p, const std::string kernel_name) : p_(p), name_(kernel_name) {}
  KernelMainCGroup(ProgramMainCGroup p,
                   const std::string kernel_name,
                   std::vector<size_t> local_dims,
                   std::vector<size_t> global_dims)
    : p_(p), name_(kernel_name) {
    // TODO: get the individual kernel function to invoke
  }

  KernelInvokeMainCGroup invoke() const { // TODO: add parameters
    return KernelInvokeMainCGroup{};
  }

public:
  ProgramMainCGroup p_;
  std::string name_;
  std::array<size_t, 3> ldims_  = {};
  std::array<size_t, 3> gdims_  = {};
  std::string kernel_module_;
  
};

} // namespace backend
} // namespace txdnn